"""
数据库管理模块 - SQLite实现
提供成本记录的存储与查询功能
"""
import sqlite3
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, Optional, List
from loguru import logger


class Database:
    """数据库管理类"""
    
    def __init__(self, db_path: str = "./database/storyforge.db"):
        """
        初始化数据库连接
        
        Args:
            db_path: 数据库文件路径
        """
        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._init_db()
    
    def _init_db(self):
        """初始化数据库表结构"""
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            
            # 创建成本记录表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS cost_records (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp TEXT NOT NULL,
                    project TEXT NOT NULL,
                    type TEXT NOT NULL,
                    provider TEXT,
                    tokens_or_items INTEGER,
                    unit_cost REAL NOT NULL,
                    total REAL NOT NULL,
                    currency TEXT DEFAULT 'USD',
                    meta TEXT
                )
            """)
            
            # 创建索引
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_cost_time 
                ON cost_records(timestamp)
            """)
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_cost_proj_type 
                ON cost_records(project, type)
            """)
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_cost_provider 
                ON cost_records(provider)
            """)
            
            conn.commit()
            logger.debug(f"数据库初始化完成: {self.db_path}")
        except Exception as e:
            logger.error(f"数据库初始化失败: {e}")
            raise
        finally:
            conn.close()
    
    def _get_connection(self) -> sqlite3.Connection:
        """获取数据库连接"""
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        return conn
    
    def insert_cost_record(
        self,
        project: str,
        cost_type: str,
        provider: str,
        tokens_or_items: int,
        unit_cost: float,
        total: float,
        currency: str = "USD",
        meta: Optional[Dict[str, Any]] = None
    ) -> int:
        """
        插入成本记录
        
        Args:
            project: 项目名称
            cost_type: 成本类型（llm/image/tts/compose等）
            provider: 提供商名称
            tokens_or_items: token数或物品数
            unit_cost: 单价
            total: 总价
            currency: 货币代码
            meta: 额外元数据（JSON）
            
        Returns:
            记录ID
        """
        import json
        
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT INTO cost_records 
                (timestamp, project, type, provider, tokens_or_items, unit_cost, total, currency, meta)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                datetime.now().isoformat(),
                project,
                cost_type,
                provider,
                tokens_or_items,
                unit_cost,
                total,
                currency,
                json.dumps(meta) if meta else None
            ))
            conn.commit()
            record_id = cursor.lastrowid
            logger.debug(f"成本记录已插入: ID={record_id}, {cost_type} {provider} ${total:.4f}")
            return record_id
        except Exception as e:
            logger.error(f"插入成本记录失败: {e}")
            conn.rollback()
            raise
        finally:
            conn.close()
    
    def get_cost_summary(
        self,
        project: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        获取成本汇总
        
        Args:
            project: 项目名称（可选）
            start_date: 开始日期（ISO-8601）
            end_date: 结束日期（ISO-8601）
            
        Returns:
            成本汇总字典
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            
            # 构建查询条件
            conditions = []
            params = []
            
            if project:
                conditions.append("project = ?")
                params.append(project)
            if start_date:
                conditions.append("timestamp >= ?")
                params.append(start_date)
            if end_date:
                conditions.append("timestamp <= ?")
                params.append(end_date)
            
            where_clause = " WHERE " + " AND ".join(conditions) if conditions else ""
            
            # 查询总成本
            cursor.execute(f"""
                SELECT SUM(total) as total_cost
                FROM cost_records
                {where_clause}
            """, params)
            row = cursor.fetchone()
            total_cost = row['total_cost'] if row['total_cost'] else 0.0
            
            # 按类型汇总
            cursor.execute(f"""
                SELECT type, SUM(total) as cost
                FROM cost_records
                {where_clause}
                GROUP BY type
            """, params)
            by_type = {row['type']: row['cost'] for row in cursor.fetchall()}
            
            # 按提供商汇总
            cursor.execute(f"""
                SELECT provider, SUM(total) as cost
                FROM cost_records
                {where_clause}
                GROUP BY provider
            """, params)
            by_provider = {row['provider']: row['cost'] for row in cursor.fetchall()}
            
            return {
                "total": total_cost,
                "by_type": by_type,
                "by_provider": by_provider
            }
        except Exception as e:
            logger.error(f"查询成本汇总失败: {e}")
            return {"total": 0.0, "by_type": {}, "by_provider": {}}
        finally:
            conn.close()
    
    def get_cost_records(
        self,
        project: Optional[str] = None,
        limit: int = 100,
        offset: int = 0
    ) -> List[Dict[str, Any]]:
        """
        获取成本记录列表
        
        Args:
            project: 项目名称（可选）
            limit: 返回记录数量
            offset: 偏移量
            
        Returns:
            成本记录列表
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            
            where_clause = "WHERE project = ?" if project else ""
            params = [project] if project else []
            
            cursor.execute(f"""
                SELECT *
                FROM cost_records
                {where_clause}
                ORDER BY timestamp DESC
                LIMIT ? OFFSET ?
            """, params + [limit, offset])
            
            records = []
            for row in cursor.fetchall():
                record = dict(row)
                # 解析meta JSON
                if record.get('meta'):
                    import json
                    try:
                        record['meta'] = json.loads(record['meta'])
                    except:
                        pass
                records.append(record)
            
            return records
        except Exception as e:
            logger.error(f"查询成本记录失败: {e}")
            return []
        finally:
            conn.close()


# 全局数据库实例
_db_instance: Optional[Database] = None


def get_database(db_path: str = "./database/storyforge.db") -> Database:
    """
    获取数据库实例（单例）
    
    Args:
        db_path: 数据库文件路径
        
    Returns:
        Database实例
    """
    global _db_instance
    if _db_instance is None:
        _db_instance = Database(db_path)
    return _db_instance


def log_cost_to_db(
    project: str,
    cost_type: str,
    provider: str,
    tokens_or_items: int,
    unit_cost: float,
    total: float,
    currency: str = "USD",
    meta: Optional[Dict[str, Any]] = None,
    db_path: str = "./database/storyforge.db"
) -> int:
    """
    记录成本到数据库
    
    Args:
        project: 项目名称
        cost_type: 成本类型
        provider: 提供商
        tokens_or_items: token数或物品数
        unit_cost: 单价
        total: 总价
        currency: 货币代码
        meta: 额外元数据
        db_path: 数据库路径
        
    Returns:
        记录ID
    """
    db = get_database(db_path)
    return db.insert_cost_record(
        project, cost_type, provider, tokens_or_items,
        unit_cost, total, currency, meta
    )


def get_cost_report(
    project: Optional[str] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    db_path: str = "./database/storyforge.db"
) -> Dict[str, Any]:
    """
    获取成本报告
    
    Args:
        project: 项目名称（可选）
        start_date: 开始日期
        end_date: 结束日期
        db_path: 数据库路径
        
    Returns:
        成本报告字典
    """
    db = get_database(db_path)
    return db.get_cost_summary(project, start_date, end_date)

